{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\nCompile Caffe2 Models\n=====================\n**Author**: `Hiroyuki Makino <https://makihiro.github.io/>`_\n\nThis article is an introductory tutorial to deploy Caffe2 models with Relay.\n\nFor us to begin with, Caffe2 should be installed.\n\nA quick solution is to install via conda\n\n.. code-block:: bash\n\n    # for cpu\n    conda install pytorch-nightly-cpu -c pytorch\n    # for gpu with CUDA 8\n    conda install pytorch-nightly cuda80 -c pytorch\n\nor please refer to official site\nhttps://caffe2.ai/docs/getting-started.html\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Load pretrained Caffe2 model\n----------------------------\nWe load a pretrained resnet50 classification model provided by Caffe2.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from caffe2.python.models.download import ModelDownloader\n\nmf = ModelDownloader()\n\n\nclass Model:\n    def __init__(self, model_name):\n        self.init_net, self.predict_net, self.value_info = mf.get_c2_model(model_name)\n\n\nresnet50 = Model(\"resnet50\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Load a test image\n------------------\nA single cat dominates the examples!\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from tvm.contrib.download import download_testdata\nfrom PIL import Image\nfrom matplotlib import pyplot as plt\nimport numpy as np\n\nimg_url = \"https://github.com/dmlc/mxnet.js/blob/main/data/cat.png?raw=true\"\nimg_path = download_testdata(img_url, \"cat.png\", module=\"data\")\nimg = Image.open(img_path).resize((224, 224))\nplt.imshow(img)\nplt.show()\n# input preprocess\ndef transform_image(image):\n    image = np.array(image) - np.array([123.0, 117.0, 104.0])\n    image /= np.array([58.395, 57.12, 57.375])\n    image = image.transpose((2, 0, 1))\n    image = image[np.newaxis, :].astype(\"float32\")\n    return image\n\n\ndata = transform_image(img)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Compile the model on Relay\n--------------------------\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Caffe2 input tensor name, shape and type\ninput_name = resnet50.predict_net.op[0].input[0]\nshape_dict = {input_name: data.shape}\ndtype_dict = {input_name: data.dtype}\n\n# parse Caffe2 model and convert into Relay computation graph\nfrom tvm import relay, transform\n\nmod, params = relay.frontend.from_caffe2(\n    resnet50.init_net, resnet50.predict_net, shape_dict, dtype_dict\n)\n\n# compile the model\n# target x86 CPU\ntarget = \"llvm\"\nwith transform.PassContext(opt_level=3):\n    lib = relay.build(mod, target, params=params)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Execute on TVM\n---------------\nThe process is no different from other examples.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import tvm\nfrom tvm import te\nfrom tvm.contrib import graph_executor\n\n# context x86 CPU, use tvm.cuda(0) if you run on GPU\ndev = tvm.cpu(0)\n# create a runtime executor module\nm = graph_executor.GraphModule(lib[\"default\"](dev))\n# set inputs\nm.set_input(input_name, tvm.nd.array(data.astype(\"float32\")))\n# execute\nm.run()\n# get outputs\ntvm_out = m.get_output(0)\ntop1_tvm = np.argmax(tvm_out.numpy()[0])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Look up synset name\n-------------------\nLook up prediction top 1 index in 1000 class synset.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from caffe2.python import workspace\n\nsynset_url = \"\".join(\n    [\n        \"https://gist.githubusercontent.com/zhreshold/\",\n        \"4d0b62f3d01426887599d4f7ede23ee5/raw/\",\n        \"596b27d23537e5a1b5751d2b0481ef172f58b539/\",\n        \"imagenet1000_clsid_to_human.txt\",\n    ]\n)\nsynset_name = \"imagenet1000_clsid_to_human.txt\"\nsynset_path = download_testdata(synset_url, synset_name, module=\"data\")\nwith open(synset_path) as f:\n    synset = eval(f.read())\nprint(\"Relay top-1 id: {}, class name: {}\".format(top1_tvm, synset[top1_tvm]))\n# confirm correctness with caffe2 output\np = workspace.Predictor(resnet50.init_net, resnet50.predict_net)\ncaffe2_out = p.run({input_name: data})\ntop1_caffe2 = np.argmax(caffe2_out)\nprint(\"Caffe2 top-1 id: {}, class name: {}\".format(top1_caffe2, synset[top1_caffe2]))"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.6.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}